#include "MeshGeometry.h"
#include "Material.h"
#include "Shader.h"
#include "RayTraceSystem.h"
#include <algorithm>
#include <map>
#include "Debug.h"
#include "ObjModel.h"
#include "Basic/FileSystem.h"

using namespace GxLibBasic;

namespace RayTracePro
{
	const int MeshKdTreeMaxDepth = 32;

	class ModelFileNotSupportedException : public Exception
	{
	public:
		ModelFileNotSupportedException(const String & fileName)
			: Exception(String(L"Can't load model \'") + fileName + L"\', no loader found for the specified model type", __FILE__, __LINE__)
		{
		}
	};

	RefPtr<Mapping<String, GeometryLoaderFunction>> GeometryLoader::loaders;

	class MeshKdTreeEvaluator
	{
	public:
		static const float IntersectTime;
		static const float TraversalTime;
		static const float EmptyNodeBounus;
		static const int ElementsPerNode = 1;
	};

	const float MeshKdTreeEvaluator::IntersectTime = 80.0f;
	const float MeshKdTreeEvaluator::TraversalTime = 1.0f;
	const float MeshKdTreeEvaluator::EmptyNodeBounus = 0.4f;

	class MeshTracingTimer
	{
		int id;
	public:
		MeshTracingTimer(int _id)
		{
			id = _id;
			DebugTimming::StartCounter(id);
		}
		~MeshTracingTimer()
		{
			DebugTimming::StopCounter(id);
		}
	};

	// Trace Triangle
	template<bool pred>
	class MeshTracer
	{
	public:
		MeshGeometry * mesh;
		inline bool Trace(DifferentialGeometry & inter, const MeshFace & face, const Ray & ray, const RayDifferential & diffIn, float & t) const
		{
			//MeshTracingTimer(pred?2:1);
			const int mod3[] = {0,1,2,0,1};
			int u = mod3[face.ProjectionAxis + 1];
			int v = mod3[face.ProjectionAxis + 2];
			float invNdotD = 1.0f / (face.PlaneU * ray.Direction[u] + face.PlaneV * ray.Direction[v] + ray.Direction[face.ProjectionAxis]);
			float tplane = -(face.PlaneU * ray.Origin[u] + face.PlaneV * ray.Origin[v] + ray.Origin[face.ProjectionAxis] + face.PlaneD) * invNdotD;
			if (tplane < ray.tMin/*-Epsilon*/ || tplane > ray.tMax/*+Epsilon*/)
				return false;
			float hitU = ray.Origin[u] + ray.Direction[u] * tplane;
			float hitV = ray.Origin[v] + ray.Direction[v] * tplane;
			
			float beta = face.K_beta_u * hitU + face.K_beta_v * hitV + face.K_beta_d;
			if (beta < -Epsilon)
				return false;
			float gamma = face.K_gamma_u * hitU + face.K_gamma_v * hitV + face.K_gamma_d;
			if (gamma < -Epsilon)
				return false;
			if (beta + gamma > 1.0f + Epsilon)
				return false;
			if (pred)
				return true;
			inter.Init();
			inter.Position[u] = hitU;
			inter.Position[v] = hitV;
			inter.Position[face.ProjectionAxis] = ray.Origin[face.ProjectionAxis] + ray.Direction[face.ProjectionAxis] * tplane;
			//float alpha = face.K_alpha_u * hitU + face.K_alpha_v * hitV + face.K_alpha_d;
			float alpha = 1.0f - beta - gamma;
			Vec3 tmp;
			Vec3::Scale(inter.Normal, mesh->Normals[face.Normal1], alpha);
			Vec3::Scale(tmp, mesh->Normals[face.Normal2], beta);
			inter.Normal += tmp;
			Vec3::Scale(tmp, mesh->Normals[face.Normal3], gamma);
			inter.Normal += tmp;

			Vec2 tmpVec2;
			if (face.TexCoord1 != -1)
			{
				inter.UV.x = mesh->TexCoords[face.TexCoord1].x * alpha;
				inter.UV.y = mesh->TexCoords[face.TexCoord1].y * alpha;
				inter.UV.x += mesh->TexCoords[face.TexCoord2].x * beta;
				inter.UV.y += mesh->TexCoords[face.TexCoord2].y * beta;
				inter.UV.x += mesh->TexCoords[face.TexCoord3].x * gamma;
				inter.UV.y += mesh->TexCoords[face.TexCoord3].y * gamma;
			}
			inter.dNdxContinuation = (void*)&face;
			t = tplane;
			if (face.MaterialId != -1)
				inter.Material = mesh->Materials[face.MaterialId].operator->();
			else
				inter.Material = mesh->Material;
			return true;
		}
	};

	MeshGeometry::MeshGeometry()
	{
		this->ComputeIntersection = MeshGeometry_ComputeIntersection;
		this->ComputeIntersectionPred = MeshGeometry_ComputeIntersectionPred;
		this->Compute_dNdx = MeshGeometry_Compute_dNdx;
		this->Bounds.Init();
	}

	void MeshGeometry::AddFace(const Vec3 verts[3], int norm[3], int tex[3], short materialId)
	{
		BBox box;
		box.Max = box.Min = verts[0];
		box.Union(verts[1]);
		box.Union(verts[2]);
		faceBBoxes.Add(box);
		Bounds.Union(box);
		MeshFace f;
		f.Normal1 = norm[0]; f.Normal2 = norm[1]; f.Normal3 = norm[2];
		f.TexCoord1 = tex[0]; f.TexCoord2 = tex[1]; f.TexCoord3 = tex[2];
		Vec3 n, c, b, a;
		Vec3::Subtract(c, verts[1], verts[0]);
		Vec3::Subtract(b, verts[2], verts[0]);
		Vec3::Subtract(a, verts[2], verts[1]);
		Vec3::Cross(n, c, b);
		Vec3::Normalize(n, n);
		int k;
		Vec3 absN = Vec3(abs(n.x), abs(n.y), abs(n.z));
		if (absN.x > absN.y && absN.x > absN.z)
			k = 0;
		else if (absN.y > absN.x && absN.y > absN.z)
			k = 1;
		else
			k = 2;
		int mod3[5] = {0, 1, 2, 0, 1};
		int u = mod3[k + 1];
		int v = mod3[k + 2];
		f.ProjectionAxis = k;
		float invNk = 1.0f / n[k];
		Vec3::Scale(n, n, invNk);
		f.PlaneU = n[u];
		f.PlaneV = n[v];
		f.PlaneD = -Vec3::Dot(n, verts[0]); 
		float divisor = 1.0f / (b[u]*c[v] - b[v]*c[u]);
		f.K_beta_u = -b[v] * divisor;
		f.K_beta_v = b[u] * divisor;
		const Vec3 & A = verts[0];
		f.K_beta_d = (b[v] * A[u] - b[u] * A[v]) * divisor;
		f.K_gamma_u = c[v] * divisor;
		f.K_gamma_v = -c[u] * divisor;
		f.K_gamma_d = (c[u] * A[v] - c[v] * A[u]) * divisor;
		const Vec3 & B = verts[1];
		f.K_alpha_u = a[v] * divisor;
		f.K_alpha_v = -a[u] * divisor;
		//f.K_alpha_d = (a[u] * B[v] - a[v]*B[u]) * divisor;
		f.MaterialId = materialId;
		Faces.Add(f);
	}

	void MeshGeometry::UpdateAccelerateStructure()
	{
		DebugTimming::Start("Creating kd-tree...");
		BBoxEdge * edges = new BBoxEdge[Faces.Count()*2];
		int * geomLeft = new int[Faces.Count()];
		int * geomRight = new int[Faces.Count()*MeshKdTreeMaxDepth];
		memset(geomLeft, 0, sizeof(int)*Faces.Count());
		memset(geomRight, 0, sizeof(int)*Faces.Count()*MeshKdTreeMaxDepth);
		for (int i = 0; i<Faces.Count(); i++)
		{
			geomLeft[i] = i;
		}
		ConstructKdTree<MeshFace, MeshKdTreeEvaluator>(tree, Faces.Buffer(), 0, faceBBoxes, 
			edges, geomLeft, Faces.Count(), geomLeft, geomRight, MeshKdTreeMaxDepth, 0);
		delete [] geomLeft;
		delete [] geomRight;
		delete [] edges;
		faceBBoxes.SetSize(0);
		faceBBoxes.Compress();
		DebugTimming::End();
	}

	void MeshGeometry::MeshGeometry_Compute_dNdx(DifferentialGeometry & inter, const Vec3 & dPdx, const Vec3 & dPdy)
	{
		const int mod3[] = {0,1,2,0,1};
		MeshFace * face = (MeshFace*)(inter.dNdxContinuation);
		MeshGeometry * mesh = (MeshGeometry*)(inter.Geometry);
		Vec3 dndx, dndy;
		Vec3 tmp, tmp2;
		int u = mod3[face->ProjectionAxis + 1];
		int v = mod3[face->ProjectionAxis + 2];
		float beta = face->K_beta_u * dPdx[u] + face->K_beta_v * dPdx[v];// + face->K_beta_d;
		float gamma = face->K_gamma_u * dPdx[u] + face->K_gamma_v * dPdx[v];// + face->K_gamma_d;
		float alpha = face->K_alpha_u * dPdx[u] + face->K_alpha_v * dPdx[v];// + face->K_alpha_d;
		Vec3::Scale(dndx, mesh->Normals[face->Normal1], alpha);
		Vec3::Scale(tmp, mesh->Normals[face->Normal2], beta);
		dndx += tmp;
		Vec3::Scale(tmp, mesh->Normals[face->Normal3], gamma);
		dndx += tmp;
		float nDotn = Vec3::Dot(inter.Normal, inter.Normal);
		float divisor = 1.0f / sqrt(nDotn * nDotn * nDotn);
		Vec3::Scale(tmp, dndx, nDotn);
		Vec3::Scale(tmp2, inter.Normal, Vec3::Dot(inter.Normal, dndx));
		tmp -= tmp2;
		Vec3::Scale(inter.dNdx, tmp, divisor);
		if (face->TexCoord1 != -1)
		{
			inter.dUdx = mesh->TexCoords[face->TexCoord1].x * alpha;
			inter.dVdx = mesh->TexCoords[face->TexCoord1].y * alpha;
			inter.dUdx += mesh->TexCoords[face->TexCoord2].x * beta;
			inter.dVdx += mesh->TexCoords[face->TexCoord2].y * beta;
			inter.dUdx += mesh->TexCoords[face->TexCoord3].x * gamma;
			inter.dVdx += mesh->TexCoords[face->TexCoord3].y * gamma;
		}
		beta = face->K_beta_u * dPdy[u] + face->K_beta_v * dPdy[v];// + face->K_beta_d;
		gamma = face->K_gamma_u * dPdy[u] + face->K_gamma_v * dPdy[v];// + face->K_gamma_d;
		alpha = face->K_alpha_u * dPdy[u] + face->K_alpha_v * dPdy[v];// + face->K_alpha_d;
		Vec3::Scale(dndy, mesh->Normals[face->Normal1], alpha);
		Vec3::Scale(tmp, mesh->Normals[face->Normal2], beta);
		dndy += tmp;
		Vec3::Scale(tmp, mesh->Normals[face->Normal3], gamma);
		dndy += tmp;
		Vec3::Scale(tmp, dndy, nDotn);
		Vec3::Scale(tmp2, inter.Normal, Vec3::Dot(inter.Normal, dndy));
		tmp -= tmp2;
		Vec3::Scale(inter.dNdy, tmp, divisor);
		if (face->TexCoord1 != -1)
		{
			inter.dUdy = mesh->TexCoords[face->TexCoord1].x * alpha;
			inter.dVdy = mesh->TexCoords[face->TexCoord1].y * alpha;
			inter.dUdy += mesh->TexCoords[face->TexCoord2].x * beta;
			inter.dVdy += mesh->TexCoords[face->TexCoord2].y * beta;
			inter.dUdy += mesh->TexCoords[face->TexCoord3].x * gamma;
			inter.dVdy += mesh->TexCoords[face->TexCoord3].y * gamma;
		}
		inter.SetDnDx();
	}
	
	bool MeshGeometry::MeshGeometry_ComputeIntersection(RayTraceSystem * system, DifferentialGeometry & inter, const Geometry * geom, 
		const Ray & ray, const RayDifferential & rayDiffIn, float & t)
	{
		MeshTracer<false> tracer;
		tracer.mesh = (MeshGeometry*)geom;
		/*float ct = FLT_MAX;
		bool rs = false;
		for (int i = 0; i<tracer.mesh->Faces.Count(); i++)
		{
			DifferentialGeometry dg;
			if (tracer.Trace(dg, tracer.mesh->Faces[i], ray, rayDiffIn, t))
			{
				if (t < ct)
				{
					ct = t;
					inter = dg;
					rs = true;
				}
			}
		}*/
		inter.dNdxContinuation = 0;
		DebugTimming::StartCounter(0);
		bool rs = TraverseKdTree<MeshFace, MeshTracer<false>, false, false>(system, tracer, inter, ((MeshGeometry*)geom)->tree, ray, rayDiffIn, t);
		DebugTimming::StopCounter(0);
		inter.Geometry = (Geometry*)geom;
		if (rs && geom->AutoFlipNormals && Vec3::Dot(inter.Normal, ray.Direction) > 0.0f)
		{
			inter.Normal.x = -inter.Normal.x;
			inter.Normal.y = -inter.Normal.y;
			inter.Normal.z = -inter.Normal.z;
		}
		return rs;
	}

	bool MeshGeometry::MeshGeometry_ComputeIntersectionPred(RayTraceSystem * system, const Geometry * geom, const Ray & ray)
	{
		MeshTracer<true> tracer;
		tracer.mesh = (MeshGeometry*)geom;
		DifferentialGeometry inter;
		RayDifferential rayDiffIn;
		float t;
		return TraverseKdTree<MeshFace, MeshTracer<true>, true, false>(system, tracer, inter, ((MeshGeometry*)geom)->tree, ray, rayDiffIn, t);
	}

	struct MeshCreateLeafBoundingSphereFunction
	{
		int bufferSize;
		Array<Vec4, true> & buffer;
		MeshCreateLeafBoundingSphereFunction(Array<Vec4, true> & _buffer)
			:buffer(_buffer)
		{}
		void operator()(BBox & curBox, int nodeId, int depth)
		{
			Vec3 position, range;
			Vec4 sphere;
			Vec3::Add(position, curBox.Min, curBox.Max);
			Vec3::Scale(position, position, 0.5f);
			sphere.x = position.x; sphere.y = position.y; sphere.z = position.z;
			Vec3::Subtract(range, curBox.Max, curBox.Min);
			sphere.w = range.Length() * 0.5f;
			buffer.Add(sphere);
		}
	};

	void MeshGeometry::ComputeBoundingSpheres(int depth, Array<Vec4, true> & buffer)
	{
		MeshCreateLeafBoundingSphereFunction cleaf(buffer);
		ComputeKdTreeBoundingSpheres(tree, depth, buffer, cleaf);
	}

	void MeshGeometry::ComputeBBox(BBox & bbox)
	{
		TransformBBox(bbox, Transformation, Bounds);
	}
	
	bool LoadObjModel(RayTraceSystem * system, MeshGeometry & geometry, const char * fileName)
	{
		ObjModel mdl;
		DebugTimming::Start((String(L"Loading ") + fileName + L"...").ToMBString());
		if (!LoadObj(mdl, fileName)) return false;
		for (int i = 0; i<mdl.Materials.Count(); i++)
		{
			Material * mat = system->CreateMaterial();
			geometry.Materials.Add(mat);
			mat->Shader = system->CreateShader(L"Lambert");
			Vec4 diffuse(mdl.Materials[i]->Diffuse, 1.0f);
			if (mdl.Materials[i]->DiffuseMap.GetLength())
			{
				String basePath = ExtractFilePath(fileName);
				String diffuseMap = basePath + mdl.Materials[i]->DiffuseMap;
				ShadingNode * texture = system->CreateShader(L"Texture");
				texture->SetParameter(L"File Name", diffuseMap.GetStrPtr());
				texture->SetParameter(L"Multiplier", diffuse);
				mat->Shader->SetParameter(L"Diffuse", texture);
				
				mat->GI.Diffuse.Set(texture);
				texture->AddRef();
			}
			else
			{
				mat->Shader->SetParameter(L"Diffuse", diffuse);
				mat->GI.Diffuse.Set(diffuse);
			}
			mat->GI.Lambert = true;
		}
		RecomputeNormals(mdl);
		geometry.Normals.Add(mdl.Normals);
		geometry.TexCoords.Add(mdl.TexCoords);
		Vec3 verts[3];
		for (int i = 0; i < mdl.Faces.Count(); i++)
		{
			ObjFace & face = mdl.Faces[i];
			verts[0] = mdl.Vertices[face.VertexIds[0]];
			verts[1] = mdl.Vertices[face.VertexIds[1]];
			verts[2] = mdl.Vertices[face.VertexIds[2]];
			geometry.AddFace(verts, face.NormalIds, face.TexCoordIds, face.MaterialId);
		}
		DebugTimming::End();
		return true;
	}
	struct FaceInFile
	{
		Vec3 Vertices[3];
		ObjFace Face;
	};
	void CreateMengerSponge(int xi, int yi, int zi, double l, int m, int level, int maxLevel, ObjModel * mdl);
	bool LoadSponge(MeshGeometry & geometry, int level)
	{
		ObjModel mdl;
		if (level == 5)
		{
			mdl.Vertices.SetCapacity(13589913*3);
			mdl.Faces.SetCapacity(13589913);
		}
		CreateMengerSponge(0, 0, 0, 1.0f, 0x3F, 0, level, &mdl);
		RecomputeNormals(mdl);
		geometry.Normals.Add(mdl.Normals);
		geometry.TexCoords.Add(mdl.TexCoords);
		Vec3 verts[3];
		int faceCount = mdl.Faces.Count();
		FILE * f;
		if (level == 5)
		{
			printf("Dumping to disk...\n");	
			fopen_s(&f, "d:\\sponge_dump.tmp", "wb");
			int countBuf = mdl.Faces.Count();
			for (int i = 0; i<countBuf; i++)
			{
				FaceInFile ff;
				ff.Face = mdl.Faces[i];
				ff.Vertices[0] = mdl.Vertices[mdl.Faces[i].VertexIds[0]];
				ff.Vertices[1] = mdl.Vertices[mdl.Faces[i].VertexIds[1]];
				ff.Vertices[2] = mdl.Vertices[mdl.Faces[i].VertexIds[2]];
				fwrite(&ff, sizeof(FaceInFile), 1, f);
			}
			fclose(f);
			fopen_s(&f, "d:\\sponge_dump.tmp", "rb");
			mdl.Faces.SetSize(0);
			mdl.Faces.Compress();
			mdl.Vertices.SetSize(0);
			mdl.Vertices.Compress();
			geometry.ReserveFaceBuffer(faceCount);
		}
		for (int i = 0; i < faceCount; i++)
		{
			ObjFace face;
			if (level == 5)
			{
				FaceInFile ff;
				fread(&ff, sizeof(FaceInFile), 1, f);
				face = ff.Face;
				verts[0] = ff.Vertices[0];
				verts[1] = ff.Vertices[1];
				verts[2] = ff.Vertices[2];
			}
			else
			{
				face = mdl.Faces[i];
				verts[0] = mdl.Vertices[face.VertexIds[0]];
				verts[1] = mdl.Vertices[face.VertexIds[1]];
				verts[2] = mdl.Vertices[face.VertexIds[2]];
			}
			geometry.AddFace(verts, face.NormalIds, face.TexCoordIds);
		}
		if (level == 5)
		{
			fclose(f);
		}
		printf("menger sponge(%d): v:%d, f:%d\n", level, mdl.Vertices.Count(), mdl.Faces.Count());
		return true;
	}
	void GeometryLoader::Init()
	{
		loaders = new Mapping<String, GeometryLoaderFunction>();
		loaders->Add(L"obj", LoadObjModel);
	}

	bool GeometryLoader::LoadGeometryFromFile(RayTraceSystem * system, MeshGeometry & geometry, const String & fileName)
	{
		GeometryLoaderFunction func;
		if (loaders->Lookup(fileName.Right(3).Lowercase(), func))
		{
			return func(system, geometry, fileName.ToMBString());
		}
		else
		{
			throw ModelFileNotSupportedException(fileName);
		}
		return false;
	}

	void CreateMengerSponge(int xi, int yi, int zi, double l, int m, int level, int maxLevel, ObjModel * mdl)
	{
		static const Vec3 msVertices[8]=
		{
			Vec3(0.0f, 1.0f, 1.0f),
			Vec3(0.0f, 0.0f, 1.0f),
			Vec3(1.0f, 0.0f, 1.0f),
			Vec3(1.0f, 1.0f, 1.0f),
			Vec3(0.0f, 1.0f, 0.0f),
			Vec3(0.0f, 0.0f, 0.0f),
			Vec3(1.0f, 0.0f, 0.0f),
			Vec3(1.0f, 1.0f, 0.0f),
		};

		static const int msTriangles[12][3]=
		{
			{0, 1, 2}, {0, 2, 3},
			{0, 4, 5}, {0, 5, 1},
			{1, 5, 6}, {1, 6, 2},
			{2, 6, 7}, {2, 7, 3},
			{3, 7, 4}, {3, 4, 0},
			{5, 4, 7}, {5, 7, 6},
		};

		static const Vec3 msNormals[12]=
		{
			Vec3(0.0f, 0.0f, 1.0f),
			Vec3(0.0f, 0.0f, 1.0f),
			Vec3(-1.0f, 0.0f, 0.0f),
			Vec3(-1.0f, 0.0f, 0.0f),
			Vec3(0.0f, -1.0f, 0.0f),
			Vec3(0.0f, -1.0f, 0.0f),
			Vec3(1.0f, 0.0f, 0.0f),
			Vec3(1.0f, 0.0f, 0.0f),
			Vec3(0.0f, 1.0f, 0.0f),
			Vec3(0.0f, 1.0f, 0.0f),
			Vec3(0.0f, 0.0f, -1.0f),
			Vec3(0.0f, 0.0f, -1.0f),
		};

		if(level==maxLevel)
		{
			const float s=2.0f/(float)l;
			Vec3 o(xi*s-1, yi*s-1, zi*s-1);
			for(int i=0;i<12;i++)
			{
				if(m&(1<<(i/2)))
				{
					Vec3 verts[3];
					Vec3::Scale(verts[0], msVertices[msTriangles[i][0]], s);
					verts[0] += o;
					Vec3::Scale(verts[1], msVertices[msTriangles[i][1]], s);
					verts[1] += o;
					Vec3::Scale(verts[2], msVertices[msTriangles[i][2]], s);
					verts[2] += o;
					Vec3::Scale(verts[0], verts[0], 1.0f);
					Vec3::Scale(verts[1], verts[1], 1.0f);
					Vec3::Scale(verts[2], verts[2], 1.0f);
					mdl->Vertices.Add(verts[0]);
					mdl->Vertices.Add(verts[1]);
					mdl->Vertices.Add(verts[2]);
					ObjFace face;
					face.VertexIds[0] = mdl->Vertices.Count()-3;
					face.VertexIds[1] = mdl->Vertices.Count()-2;
					face.VertexIds[2] = mdl->Vertices.Count()-1;
					face.SmoothGroup = 0;
					face.NormalIds[0] = -1; face.NormalIds[1] = -1; face.NormalIds[2] = -1;
					face.TexCoordIds[0] = -1; face.TexCoordIds[1] = -1; face.TexCoordIds[2] = -1;
					mdl->Faces.Add(face);
				}
			}
		}
		else
		{
			for(int i=0;i<=2;i++)for(int j=0;j<=2;j++)for(int k=0;k<=2;k++)
			{
				if((i&1)+(j&1)+(k&1)<=1)
				{
					int m2=0;

					if(i==0)m2|=(1<<1)&m;
					if(i==2)m2|=(1<<3)&m;
					if(j==0)m2|=(1<<2)&m;
					if(j==2)m2|=(1<<4)&m;
					if(k==0)m2|=(1<<5)&m;
					if(k==2)m2|=(1<<0)&m;
					
					if(i==2&&(j==1 || k==1))m2|=(1<<1);
					if(i==0&&(j==1 || k==1))m2|=(1<<3);
					if(j==2&&(k==1 || i==1))m2|=(1<<2);
					if(j==0&&(k==1 || i==1))m2|=(1<<4);
					if(k==2&&(i==1 || j==1))m2|=(1<<5);
					if(k==0&&(i==1 || j==1))m2|=(1<<0);

					CreateMengerSponge(xi*3+i, yi*3+j, zi*3+k, l*3, m2, level+1, maxLevel, mdl);
				}
			}
		}
	}
}
